import numpy as np

import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, Button, TextBox

from model.manipulator import Manipulator


# Create the robot instance
manipulator = Manipulator()

# Initialize the figure
fig, ax = plt.subplots()
plt.subplots_adjust(left=0.15, right=0.85, bottom=0.45)
ax.margins(x=0)

# Set up the sliders
axcolor = 'lightgoldenrodyellow'
axt1 = plt.axes([0.15, 0.30, 0.65, 0.03], facecolor=axcolor)
axt2 = plt.axes([0.15, 0.25, 0.65, 0.03], facecolor=axcolor)
axtb1 = plt.axes([0.82, 0.30, 0.1, 0.03], facecolor=axcolor)
axtb2 = plt.axes([0.82, 0.25, 0.1, 0.03], facecolor=axcolor)

st1 = Slider(axt1, 't1', 0, np.math.pi, valinit=np.math.pi/2)
st2 = Slider(axt2, 't2', -np.math.pi, np.math.pi, valinit=0)
st1.valtext.set_visible(False)
st2.valtext.set_visible(False)

tb1 = TextBox(axtb1, "")
tb2 = TextBox(axtb2, "")

# Set up the reset button
resetax = plt.axes([0.8, 0.15, 0.1, 0.04])
button = Button(resetax, 'Reset', color=axcolor, hovercolor='0.975')

# Global variables to store the angle values
t1, t2 = [0, 0]


def draw_robot(manipulator, theta1, theta2):
    '''
    Draw the robot given the robot model and joint values

    manipulator:  The robot instance (class Manipulator)
    theta1: The value of joint1
    theta2: The value of joint2
    '''
    global ax

    p_l1 = manipulator.get_link1_pos(theta1, theta2)
    p_ee = manipulator.get_ee_pos(theta1, theta2)

    # Clear the current axes
    ax.clear()

    # Show the robot angles on the title
    text = r'$\theta_1$={0:.3f} | $\theta_2$={1:.3f}'.format(theta1, theta2)
    ax.set_title(text)

    # Draw the robot
    ax.plot([0, p_l1[0]], [0, p_l1[1]], color='blue', linewidth=4, zorder=1)
    ax.plot([p_l1[0], p_ee[0]], [p_l1[1], p_ee[1]], color='lightblue', linewidth=4, zorder=1)
    # Draw ground
    ax.plot([-5, 5], [0, 0], color='gray', linewidth=1, zorder=1)
    # Draw base
    ax.plot([-0.1, 0, 0.1, -0.1], [0, 0.173, 0, 0], color='gray', linewidth=4, zorder=2)
    # Draw points
    j1 = plt.Circle((p_l1[0], p_l1[1]), 0.16, fill=True, zorder=3)
    j2 = plt.Circle((p_ee[0], p_ee[1]), 0.12, fill=True, zorder=3)
    ax.add_artist(j1)
    ax.add_artist(j2)

    # Set the limit to x-axis and y-axis
    ax.set_xlim(-4, 4)
    ax.set_ylim(-2, 6)
    ax.set_aspect('equal', adjustable='box')

    plt.draw()

def update(val):
    '''
    Event handler for slider st1, st2
    '''
    global t1, t2, manipulator

    t1 = st1.val
    t2 = st2.val

    tb1.set_val('{:.3f}'.format(t1))
    tb2.set_val('{:.3f}'.format(t2))

    draw_robot(manipulator, t1, t2)
# Register the event handler to the sliders
st1.on_changed(update)
st2.on_changed(update)


def reset(event):
    '''
    Event handler for reset button
    '''
    st1.reset()
    st2.reset()
# Register the event handler to the reset button
button.on_clicked(reset)


def submit_t1(expression):
    '''
    Event handler for textbox1
    '''
    global t1, t2, manipulator

    # write your codes here...

def submit_t2(expression):
    '''
    Event handler for textbox2
    '''
    global t1, t2, manipulator

    # write your codes below:
    # t2 = eval(expression) # get the value of the textbox and save it to t2
    # update the value of the slider
    # update the pose of the robot

tb1.on_submit(submit_t1)
tb2.on_submit(submit_t2)

update(0)
plt.show()
